from typing import Tuple

import playwright
import pytest

from browsergym.core.env import BrowserEnv
from browsergym.core.task import AbstractBrowserTask


class MockImageGoalTask(AbstractBrowserTask):
    @classmethod
    def get_task_id(cls):
        return "mockimagegoal"

    def __init__(self, seed: int = 0, start_url: str = "https://www.google.com") -> None:
        """
        Args:
            seed: random seed.
            start_url: str, the url for the starting page.
            goal: str, the initial goal.

        """
        super().__init__(seed)
        self.start_url = start_url
        self.goal = [
            {"type": "text", "text": "This is a mock task with an image goal."},
            {
                "type": "image_url",
                "image_url": "",
            },
        ]

    def setup(self, page: playwright.sync_api.Page) -> tuple[str, dict]:
        page.goto(self.start_url, timeout=10000)
        return self.goal, {}

    def teardown(self) -> None:
        pass

    def validate(
        self, page: playwright.sync_api.Page, chat_messages: list[str]
    ) -> Tuple[float, bool, str, dict]:
        reward, done, msg, info = 0, False, "", {}

        for message in chat_messages:
            if message["role"] == "user" and message["message"] == "exit":
                done = True
                break

        return reward, done, msg, info


def test_mock_image_goal_task():
    env = BrowserEnv(MockImageGoalTask)
    obs, _ = env.reset()

    assert "goal_object" in obs
    assert len(obs["goal_object"]) == 2
    assert obs["goal_object"][0]["type"] == "text"
    assert obs["goal_object"][0]["text"] == "This is a mock task with an image goal."
    assert obs["goal_object"][1]["type"] == "image_url"

    env.chat.add_message("user", "exit")
    obs, reward, terminated, _, _ = env.step("send_msg_to_user('bye')")

    assert reward == 0
    assert terminated is True

    env.close()


if __name__ == "__main__":
    test_mock_image_goal_task()
